package com.sz.classload;

import org.jetbrains.annotations.NotNull;

import java.util.concurrent.*;

public class PriorityThreadPoolExecutor extends ThreadPoolExecutor {
    private ThreadLocal<Integer> threadLocal = new ThreadLocal<Integer>(){
        @Override
        protected Integer initialValue() {
            return 0;
        }
    };

    public PriorityThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, getWorkQueue());
    }

    public PriorityThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit,  getWorkQueue(), threadFactory);
    }

    public PriorityThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit,  @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit,  getWorkQueue(), handler);
    }

    public PriorityThreadPoolExecutor(int corePoolSize, int maximumPoolSize, long keepAliveTime, @NotNull TimeUnit unit, @NotNull ThreadFactory threadFactory, @NotNull RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit,  getWorkQueue(), threadFactory, handler);
    }

    public static PriorityBlockingQueue getWorkQueue(){
        return new PriorityBlockingQueue();
    }

    @Override
    public void execute(@NotNull Runnable command) {
        int priority = threadLocal.get();
        try{
            this.execute(command,priority);
        }finally {
            threadLocal.set(0);
        }
    }

    public void execute(Runnable runnable,int priority){
        super.execute(new PriorityRunnable(runnable,priority));
    }

    public Future<?> submit(@NotNull Runnable task,int priority) {
        threadLocal.set(priority);
        return super.submit(task);
    }

    public <T> Future<T> submit(@NotNull Callable<T> task,int pririty) {
        threadLocal.set(pririty);
        return super.submit(task);
    }

    public <T> Future<T> submit(@NotNull Runnable task, T result,int priority) {
        threadLocal.set(priority);
        return super.submit(task, result);
    }
}
